package cn.caplike.demo.repository.javaagent.agentmain;


import javassist.*;

import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.reflect.Proxy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.stream.Collectors;

@SuppressWarnings({"rawtypes", "unchecked"})
public class Agent {

    // todo: 利用 java 的事件机制 attach

    /**
     * @param instrumentation {@link Instrumentation}
     */
    public static void agentmain(String agentArgs, Instrumentation instrumentation) {
        final TheClassFileTransformer transformer = new TheClassFileTransformer();
        instrumentation.addTransformer(transformer, true);

        System.out.println("Agentmain running.");
        Arrays.stream(instrumentation.getAllLoadedClasses())
                .filter(clazz -> {
                    try {
                        return transformer.supported(clazz);
                    } catch (Throwable ignored) {
                        System.err.println(ignored.getMessage());
                        return false;
                    }
                })
                .forEach(the -> {
                            try {
                                instrumentation.retransformClasses(the);
                                // instrumentation.redefineClasses(new ClassDefinition(the, redefine(the)));
                            } catch (UnmodifiableClassException e) {
                                throw new RuntimeException(e);
                            } /*catch (ClassNotFoundException e) {
                                throw new RuntimeException(e);
                            }*/
                        }
                );
    }

    public static byte[] redefine(Class<?> classBeingRedefined) {
        final ClassPool classPool = ClassPool.getDefault();
        try {
            final CtClass ctClass = classPool.get(classBeingRedefined.getName());
            CtMethod setter = ctClass.getDeclaredMethod("setName", classPool.get(new String[]{"java.lang.String"}));
            setter.insertBefore("System.out.println(\"This method has been modified.\");");
            setter.insertAfter("super.put(\"name\", name);");

            try {
//                final CtMethod newGenerated = ctClass.getDeclaredMethod("newGenerated");
//                System.err.println(newGenerated.getName());

                //新增一个方法: javaagent 不能
                CtMethod ctMethod = new CtMethod(CtClass.voidType, "newGenerated", new CtClass[]{}, ctClass);
                ctMethod.setModifiers(Modifier.PUBLIC);
                ctMethod.setBody("{System.out.println(\"This method is totally new.\");}");
                ctClass.addMethod(ctMethod);

                ctClass.writeFile(Agent.class.getProtectionDomain().getCodeSource().getLocation().getPath());
                return ctClass.toBytecode();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        } catch (NotFoundException | CannotCompileException e) {
            throw new RuntimeException(e);
        }
    }
}
